"""
任务管理器 - 处理A2A协议的任务请求
"""
import asyncio
import logging
from abc import ABC, abstractmethod
from typing import Any, Dict, Optional, AsyncGenerator
from uuid import uuid4
from datetime import datetime

from .types import (
    Task, TaskStatus, TaskState, Message, TextPart, DataPart, FilePart, Artifact,
    TaskSendParams, TaskQueryParams, TaskNotFoundError,
    TaskStatusUpdateEvent, TaskArtifactUpdateEvent
)

logger = logging.getLogger(__name__)


def extract_text_from_part(part) -> str:
    """从Part对象中安全地提取文本内容"""
    # 处理 RootModel 结构
    if hasattr(part, 'root'):
        actual_part = part.root
    else:
        actual_part = part
    
    # 检查是否为 TextPart 或有 text 属性
    if isinstance(actual_part, TextPart):
        return actual_part.text
    elif hasattr(actual_part, 'text'):
        return actual_part.text
    elif hasattr(actual_part, 'kind') and actual_part.kind == 'text' and hasattr(actual_part, 'text'):
        return actual_part.text
    else:
        return str(actual_part)


class TaskManager(ABC):
    """任务管理器抽象基类"""
    
    @abstractmethod
    async def handle_send_task(self, params: Dict[str, Any]) -> Task:
        """处理发送任务请求"""
        pass
    
    @abstractmethod
    async def handle_get_task(self, params: Dict[str, Any]) -> Task:
        """处理获取任务请求"""
        pass

    @abstractmethod
    async def create_task(self, params: TaskSendParams) -> Task:
        """创建任务"""
        pass

    @abstractmethod
    async def execute_task_streaming(self, task: Task, message: Message) -> AsyncGenerator[TaskStatusUpdateEvent | TaskArtifactUpdateEvent, None]:
        """处理获取任务请求"""
        pass

    @abstractmethod
    async def emit_status_event(self, task: Task, state: TaskState, event_message: str | Dict[str, Any],
                                metadata: Optional[Dict[str, Any]] = None,
                                final: bool = False) -> TaskStatusUpdateEvent | TaskArtifactUpdateEvent:
        pass

class InMemoryTaskManager(TaskManager):
    """内存中的任务管理器"""
    
    def __init__(self):
        self.tasks: Dict[str, Task] = {}
        self.lock = asyncio.Lock()
        # 添加事件监听器字典，用于存储流式事件处理函数
        self.event_listeners: Dict[str, Any] = {}

    def add_event_listener(self, task_id: str, listener):
        """添加事件监听器"""
        self.event_listeners[task_id] = listener

    def remove_event_listener(self, task_id: str):
        """移除事件监听器"""
        if task_id in self.event_listeners:
            del self.event_listeners[task_id]

    async def _emit_status_event(self, task: Task, state: TaskState, event_message: str | Dict[str, Any], metadata: Optional[Dict[str, Any]] = None, final: bool = False) -> TaskStatusUpdateEvent | TaskArtifactUpdateEvent:
        """
            发送任务状态更新事件 - 符合A2A协议标准

            Args:
                task: 任务对象
                state: 任务状态
                event_message: 事件消息
                                    - 可以是字符串或字典
                                    - 如果是工件、文件或者数据，则为字典格式，
                                    - 其中工件包含 `type` = 'artifact', `name`, `description`, `text`, `append`, `lastChunk` 等字段，artifactType 字段可选值['report'、'image'、'dataset'、'mermaid']等
                                    - 其中文件包含 `file` = {"name": "filename", "mimeType": "image/png", "bytes":"base64_data"} / {"name": "filename", "mimeType": "text/plain", "uri":"文件下载uri"}字段
                                    - 其中数据包含 `data` 字段，值为字典、列表、字符串等
                                        - 例如iframe数据可以是 {"data": {"iframe": "IframeData类型数据", "dataType": "iframe"}}
                metadata: 元数据
                                - 可选，包含额外的事件信息
                                - 如果是工件，则可以包含 `dataType` 和 `artifactType` 字段，`dataType` = 'artifact' 和 `artifactType`可选值['report'、'image'、'dataset'、'mermaid']
                final: 是否为最终事件
        """
        context_id = getattr(task, 'sessionId', None) or task.id

        is_artifact = False
        message = None
        if isinstance(event_message, str):
            message = Message(
                messageId=str(uuid4()),
                role="agent",
                parts=[TextPart(text=event_message, metadata=metadata)] if event_message else [],
                contextId=context_id
            ) if event_message else None
        elif isinstance(event_message, dict):
            if event_message.get('type', 'message') == 'artifact':
                is_artifact = True
                if metadata is None:
                    metadata = {}
                metadata["dataType"] = "artifact"
                if "artifactType" in event_message:
                    metadata["artifactType"] = event_message.get('artifactType', 'mermaid')
                elif "artifactType" not in metadata:
                    metadata["artifactType"] = "mermaid"

                message = Artifact(
                    name=event_message.get('name', None),
                    description=event_message.get('description', None),
                    parts=[TextPart(
                        kind="text",
                        text=event_message.get('text', '')
                    )],
                    metadata=metadata,
                    index=0,
                    append=event_message.get('append', True),
                    lastChunk=event_message.get('lastChunk', True),
                )
            else:
                if "file" in event_message:
                    file_part = FilePart(
                        file=event_message['file'],
                        metadata=metadata
                    )
                    message = Message(
                        messageId=str(uuid4()),
                        role="agent",
                        parts=[file_part],
                        contextId=context_id
                    )
                elif "data" in event_message:
                    data_part = DataPart(
                        data=event_message['data'],
                        metadata=metadata
                    )
                    message = Message(
                        messageId=str(uuid4()),
                        role="agent",
                        parts=[data_part],
                        contextId=context_id
                    )
                elif "text" in event_message:
                    message = Message(
                        messageId=str(uuid4()),
                        role="agent",
                        parts=[TextPart(text=event_message.get('text'), metadata=metadata)],
                        contextId=context_id
                    )
                else:
                    message = None

        if is_artifact:
            status_event = TaskArtifactUpdateEvent(
                append=False,
                contextId=task.contextId,
                kind='artifact-update',
                taskId=task.id,
                lastChunk=True,
                artifact=message,
            )
            message = None
        else:
            status_event = TaskStatusUpdateEvent(
                taskId=task.id,
                contextId=context_id,
                status=TaskStatus(
                    state=state,
                    message=message,
                    timestamp=datetime.now()
                ),
                final=final,
                metadata=task.metadata
            )

        # 如果有事件监听器，则发送事件
        if task.id in self.event_listeners:
            listener = self.event_listeners[task.id]
            if hasattr(listener, '__call__'):
                try:
                    await listener(status_event)
                except Exception as e:
                    logger.error(f"事件监听器处理错误: {e}")

        # 更新任务状态
        await self._update_task_status(
            task.id,
            state,
            message
        )

        return status_event

    async def emit_status_event(self, task: Task, state: TaskState, event_message: str | Dict[str, Any], metadata: Optional[Dict[str, Any]] = None, final: bool = False) -> TaskStatusUpdateEvent | TaskArtifactUpdateEvent:
        """
            发送任务状态更新事件 - 符合A2A协议标准

            Args:
                task: 任务对象
                state: 任务状态
                event_message: 事件消息
                                    - 可以是字符串或字典
                                    - 如果是工件、文件或者数据，则为字典格式，
                                    - 其中工件包含 `type` = 'artifact', `name`, `description`, `text`, `append`, `lastChunk` 等字段，artifactType 字段可选值['report'、'image'、'dataset'、'mermaid']等
                                    - 其中文件包含 `file` = {"name": "filename", "mimeType": "image/png", "bytes":"base64_data"} / {"name": "filename", "mimeType": "text/plain", "uri":"文件下载uri"}字段
                                    - 其中数据包含 `data` 字段，值为字典、列表、字符串等
                                        - 例如iframe数据可以是 {"data": {"iframe": "IframeData类型数据", "dataType": "iframe"}}
                metadata: 元数据
                                - 可选，包含额外的事件信息
                                - 如果是工件，则可以包含 `dataType` 和 `artifactType` 字段，`dataType` = 'artifact' 和 `artifactType`可选值['report'、'image'、'dataset'、'mermaid']
                final: 是否为最终事件
        """
        return await self._emit_status_event(task, state, event_message, metadata, final)

    async def _emit_status_event_with_message(self, task: Task, state: TaskState, message: Message,
                                              final: bool = False):
        """发送包含完整消息的状态事件"""
        context_id = task.contextId or ""

        # 创建任务状态
        task_status = TaskStatus(
            state=state,
            message=message,
            timestamp=datetime.now()
        )

        # 创建状态更新事件
        status_event = TaskStatusUpdateEvent(
            taskId=task.id,
            contextId=context_id,
            kind="status-update",
            status=task_status,
            final=final
        )

        # 如果有事件监听器，则发送事件
        if task.id in self.event_listeners:
            listener = self.event_listeners[task.id]
            if hasattr(listener, '__call__'):
                try:
                    await listener(status_event)
                except Exception as e:
                    logger.error(f"事件监听器处理错误: {e}")

        # 直接更新任务状态，不通过_update_task_status避免重复事件
        async with self.lock:
            if task.id in self.tasks:
                self.tasks[task.id].status = task_status
    async def _update_task_status(self, task_id: str, state: TaskState, message: Optional[Message] = None):
        """更新任务状态"""
        async with self.lock:
            task = self.tasks.get(task_id)
            if task:
                task.status = TaskStatus(
                    state=state,
                    message=message,
                    timestamp=datetime.now()
                )

    async def handle_send_task(self, params: Dict[str, Any]) -> Task:
        """处理发送任务请求"""
        try:
            # 解析参数
            task_params = TaskSendParams(**params)
            
            # 创建或更新任务
            task = await self.create_task(task_params)
            
            # 执行任务
            await self._execute_task(task, task_params.message)
            
            return task
            
        except Exception as e:
            logger.error(f"处理发送任务请求时发生错误: {e}")
            raise
    
    async def handle_get_task(self, params: Dict[str, Any]) -> Task:
        """处理获取任务请求"""
        try:
            query_params = TaskQueryParams(**params)
            
            async with self.lock:
                task = self.tasks.get(query_params.id)
                if task is None:
                    raise TaskNotFoundError()
                
                return task
                
        except Exception as e:
            logger.error(f"处理获取任务请求时发生错误: {e}")
            raise
    
    async def create_task(self, params: TaskSendParams) -> Task:
        """创建任务"""
        async with self.lock:
            # 生成contextId，用于跨交互的上下文对齐
            context_id = params.sessionId or str(uuid4())
            
            task = Task(
                id=params.id,
                contextId=context_id,
                status=TaskStatus(
                    state=TaskState.SUBMITTED,
                    timestamp=datetime.now()
                ),
                history=[params.message],  # 初始化消息历史
                artifacts=None,
                metadata=params.metadata,
                kind='task',
                sessionId=params.sessionId  # 保持向后兼容
            )
            
            self.tasks[task.id] = task
            return task
    
    async def _execute_task(self, task: Task, message: Message):
        """执行任务 - 子类需要重写此方法"""
        logger.info(f"执行任务: {task.id}")
        # 更新任务状态为工作中
        await self.emit_status_event(task, TaskState.WORKING, "开始执行任务")
        await asyncio.sleep(0.001)  # 极短暂的延迟确保事件被处理

        # 默认实现：简单回复
        logger.info(f"反馈任务: {task.id}")
        response_text = f"收到消息: {extract_text_from_part(message.parts[0]) if message.parts else '无内容'}"
        await self.emit_status_event(task, TaskState.WORKING, response_text)
        await asyncio.sleep(0.001)  # 极短暂的延迟确保事件被处理

        response_message = Message(
            messageId=str(uuid4()),
            role="agent",
            parts=[TextPart(text=response_text)],
            contextId=task.contextId,
            taskId=task.id,
            kind='message'
        )
        
        # 添加AI回复到历史记录
        async with self.lock:
            if task.history is None:
                task.history = []
            task.history.append(response_message)

        logger.info(f"结束任务: {task.id}")
        # 更新任务状态为完成，并设置最终的响应消息
        #await self._update_task_status(task.id, TaskState.COMPLETED, response_message)
        # 发送任务完成事件 - 包含最终消息
        await self.emit_status_event(
            task,
            TaskState.COMPLETED,
            "",  # 最终事件不包含累积文本，只表示任务完成
            final=True
        )
        await asyncio.sleep(0.001)  # 极短暂的延迟确保事件被处理

    async def execute_task_streaming(self, task: Task, message: Message) -> AsyncGenerator[TaskStatusUpdateEvent | TaskArtifactUpdateEvent, None]:
        """流式执行任务 - 返回事件生成器"""
        event_queue = asyncio.Queue()

        # 设置事件监听器
        async def event_handler(update_event: TaskStatusUpdateEvent | TaskArtifactUpdateEvent):
            await event_queue.put(update_event)
            # 只保留简单的调试信息，不显示详细内容
            # logger.debug(f"Event queued: {update_event.kind}")

        self.add_event_listener(task.id, event_handler)

        is_final = False
        try:
            # 在后台启动任务执行
            task_execution = asyncio.create_task(self._execute_task(task, message))

            # 流式返回事件
            while True:
                try:
                    # 等待事件，设置较短的超时以提高响应性
                    event = await asyncio.wait_for(event_queue.get(), timeout=1.0)

                    # 如果是最终事件，结束流式返回
                    if isinstance(event, TaskStatusUpdateEvent) and event.final:
                        logger.debug("Final event received, ending stream")
                        is_final = True

                    yield event

                    if is_final:
                        break

                except asyncio.TimeoutError:
                    # 检查任务是否仍在执行
                    if task_execution.done():
                        # 任务已完成但没有收到最终事件，可能出现了问题
                        logger.warning(f"任务 {task.id} 已完成但未收到最终事件")
                        break
                    # 任务仍在执行，继续等待
                    continue

        except Exception as e:
            logger.error(f"流式执行过程中发生错误: {e}")
            is_final = True
            # 发送错误事件
            error_event = await self.emit_status_event(
                task,
                TaskState.FAILED,
                f"流式执行错误: {str(e)}",
                final=True
            )
            yield error_event
        finally:
            if not is_final:
                final_event = await self.emit_status_event(
                    task,
                    TaskState.COMPLETED,
                    {},
                    final=True
                )
                yield final_event

            #确保后台任务已完成
            if task_execution and not task_execution.done():
                task_execution.cancel()
                try:
                    await task_execution
                except asyncio.CancelledError:
                    pass
                except Exception as e:
                    logger.error(f"取消任务 {task.id} 执行时发生错误: {e}")
            # 清理事件监听器
            self.remove_event_listener(task.id)
            logger.info(f"已清理任务 {task.id} 的事件监听器")

            while not event_queue.empty():
                try:
                    event = event_queue.get_nowait()
                    logger.debug(f"Cleaning up event: {event.kind}")
                except asyncio.QueueEmpty:
                    pass

class AITaskManager(InMemoryTaskManager):
    """基于AI的任务管理器 - 支持流式返回"""
    
    def __init__(self, ai_client, system_prompt: str = "", model: str = "gpt-4o-mini"):
        super().__init__()
        self.ai_client = ai_client
        self.system_prompt = system_prompt
        self.model = model

    async def _execute_task(self, task: Task, message: Message):
        """使用AI执行任务 - 流式返回方式"""
        try:
            # 提取用户消息文本
            user_text = ""
            for part in message.parts:
                user_text += extract_text_from_part(part) + " "
            
            user_text = user_text.strip()
            logger.info(f"处理用户消息: {user_text}")
            
            # 调用AI模型进行流式处理（内部会发送中间状态事件）
            ai_response = await self._call_ai_model_streaming(task, user_text)
            
            # 创建最终回复消息
            final_message = Message(
                messageId=str(uuid4()),
                role="agent",
                parts=[TextPart(text=ai_response)] if ai_response else [],
                contextId=task.contextId,
                taskId=task.id,
                kind='message'
            )
            
            # 添加AI回复到历史记录
            async with self.lock:
                if task.history is None:
                    task.history = []
                task.history.append(final_message)
            
            # 发送任务完成事件 - 包含最终消息
            await self.emit_status_event(
                task,
                TaskState.COMPLETED,
                "",  # 最终事件不包含累积文本，只表示任务完成
                final=True
            )

            await asyncio.sleep(0.001)  # 极短暂的延迟确保事件被处理
        except Exception as e:
            logger.error(f"执行AI任务时发生错误: {e}")
            # 创建错误消息
            error_message = Message(
                messageId=str(uuid4()),
                role="agent",
                parts=[TextPart(text=f"处理请求时发生错误: {str(e)}")],
                contextId=task.contextId,
                taskId=task.id,
                kind='message'
            )
            
            # 添加错误消息到历史记录
            async with self.lock:
                if task.history is None:
                    task.history = []
                task.history.append(error_message)
            
            # 发送任务失败事件
            await self.emit_status_event(
                task,
                TaskState.FAILED,
                f"处理请求时发生错误: {str(e)}",
                final=True
            )
    
    async def _call_ai_model_streaming(self, task: Task, user_text: str) -> str:
        """调用AI模型进行流式处理"""
        messages = []
        
        if self.system_prompt:
            messages.append({
                "role": "system",
                "content": self.system_prompt
            })
        
        messages.append({
            "role": "user", 
            "content": user_text
        })
        
        # 检查AI客户端是否支持流式处理
        if hasattr(self.ai_client.chat.completions, 'create') and hasattr(self.ai_client.chat.completions.create, '__call__'):
            try:
                # 尝试流式调用
                stream = await self.ai_client.chat.completions.create(
                    model=self.model,
                    messages=messages,
                    max_tokens=2000,
                    temperature=0.7,
                    stream=True  # 启用流式处理
                )
                
                full_response = ""
                chunk_count = 0
                
                # 处理流式响应
                async for chunk in stream:
                    if hasattr(chunk, 'choices') and len(chunk.choices) > 0:
                        choice = chunk.choices[0]
                        if hasattr(choice, 'delta') and hasattr(choice.delta, 'content'):
                            logger.info(f"处理流式响应块: {choice.delta.content}")
                            content = choice.delta.content
                            if content:
                                full_response += content
                                chunk_count += 1
                                
                                # 立即发送每个新增的内容块
                                await self.emit_status_event(
                                    task,
                                    TaskState.WORKING,
                                    content,  # 直接发送当前块的内容
                                    final=False
                                )
                                # 强制刷新，确保事件立即发送
                                await asyncio.sleep(0.001)  # 极短暂的延迟确保事件被处理
                
                return full_response.strip() if full_response else "抱歉，未能生成有效回复。"
                
            except Exception as stream_error:
                logger.warning(f"流式调用失败，回退到标准调用: {stream_error}")
                # 流式调用失败，回退到标准调用
                return await self._call_ai_model_fallback(user_text)
        else:
            # AI客户端不支持流式处理，使用标准调用
            return await self._call_ai_model_fallback(user_text)
    
    async def _call_ai_model_fallback(self, user_text: str) -> str:
        """AI模型回退调用方法"""
        messages = []
        
        if self.system_prompt:
            messages.append({
                "role": "system",
                "content": self.system_prompt
            })
        
        messages.append({
            "role": "user", 
            "content": user_text
        })
        
        # 调用AI客户端（非流式）
        response = await self.ai_client.chat.completions.create(
            model=self.model,
            messages=messages,
            max_tokens=2000,
            temperature=0.7
        )
        
        return response.choices[0].message.content.strip()
    



# 使用示例
async def streaming_task_example():
    """流式任务管理器使用示例"""
    # 模拟AI客户端
    class MockAIClient:
        class Chat:
            class Completions:
                async def create(self, **kwargs):
                    # 模拟响应
                    class MockResponse:
                        class Choice:
                            class Message:
                                content = "这是AI的回复内容"
                        choices = [Choice()]
                    return MockResponse()
        
        chat = Chat()
    
    # 创建流式任务管理器
    ai_client = MockAIClient()
    task_manager = AITaskManager(
        ai_client=ai_client,
        system_prompt="你是一个有用的AI助手",
        model="gpt-4o-mini"
    )
    
    # 创建任务
    task_params = TaskSendParams(
        id=str(uuid4()),
        message=Message(
            messageId=str(uuid4()),
            role="user",
            parts=[TextPart(text="你好，请介绍一下自己")]
        )
    )
    
    task = await task_manager._create_task(task_params)
    
    # 流式执行任务并处理事件
    print("开始流式任务执行...")
    async for event in task_manager.execute_task_streaming(task, task_params.message):
        print(f"收到事件: {event.kind}, 状态: {event.status.state}, 最终: {event.final}")
        if event.status.message and event.status.message.parts:
            message_text = extract_text_from_part(event.status.message.parts[0])
            print(f"消息内容: {message_text}")
        
        if event.final:
            print("任务执行完成")
            break


if __name__ == "__main__":
    asyncio.run(streaming_task_example()) 